package pbadger

import (
	"gitee.com/wangyubin/gutils/parsers"
	badger "github.com/dgraph-io/badger/v2"
)

type PBadgerWriter struct {
	DB           *badger.DB
	SeqBandwidth int
}

func NewPBadgerWriter(dir string) (*PBadgerWriter, error) {
	var err error
	var pb = PBadgerWriter{}

	pb.DB, err = badger.Open(badger.DefaultOptions(dir))
	if err != nil {
		pb.Close()
		return nil, err
	}

	pb.SeqBandwidth = 1000 // default 1000
	return &pb, nil
}

func (pb *PBadgerWriter) Close() {
	if pb.DB != nil {
		pb.DB.Close()
	}

	pb.DB = nil
}

func (pb *PBadgerWriter) Write(prefix string, dataCh chan []byte) error {

	seq, err := pb.DB.GetSequence([]byte(prefix), 1000)
	if err != nil {
		return err

	}
	defer seq.Release()

	txn := pb.DB.NewTransaction(true)
	for data := range dataCh {
		key, _ := seq.Next()
		if err := txn.Set(append([]byte(prefix), parsers.Uint64ToBytes(key)...), data); err == badger.ErrTxnTooBig {
			_ = txn.Commit()
			txn = pb.DB.NewTransaction(true)
			_ = txn.Set(append([]byte(prefix), parsers.Uint64ToBytes(key)...), data)
		}

	}
	return txn.Commit()

}
